import asyncio
import types
from hysdata import HSClient, api
from hysdata.store import normalize_code, Store
import pandas as pd
from tqdm.auto import tqdm

client = HSClient(auto_wait=True, speed_limit_per_minute=120)
store = Store()


# 异步API调用方式
async def float_shareholder_top10():
    df_stock = await client.async_query(api.stock_list())
    codes = df_stock['secu_code'].to_list()
    dfs = []
    for i in tqdm(range(0, len(codes)), leave=True, desc="float_shareholder_top10"):
        await client.async_auto_wait(api.float_shareholder_top10())
        df = await client.async_query(api.float_shareholder_top10(secu_code=codes[i]), cn_column_names=True,
                                      change_types=True)
        if df.shape[0] > 0:
            dfs.append(df)
        if i % 120 == 0 or i == len(codes) - 1:
            df = pd.concat(dfs)
            df['股票代码'] = normalize_code(df['证券代码'])
            df = df.drop_duplicates(subset=['股票代码', '交易日期', '股东序号'])
            await store.async_save_df(df, table_name='float_shareholder_top10',
                                      unique_cols=['股票代码', '交易日期', '股东序号'])
            dfs = []


async def shareholder_top10():
    df_stock = await client.async_query(api.stock_list())
    codes = df_stock['secu_code'].to_list()
    dfs = []
    for i in tqdm(range(0, len(codes)), leave=True, desc="shareholder_top10"):
        await client.async_auto_wait(api.shareholder_top10())
        df = await client.async_query(api.shareholder_top10(secu_code=codes[i]), cn_column_names=True,
                                      change_types=True)
        if df.shape[0] > 0:
            dfs.append(df)
        if i % 120 == 0 or i == len(codes) - 1:
            df = pd.concat(dfs)
            df['股票代码'] = normalize_code(df['证券代码'])
            df = df.drop_duplicates(subset=['股票代码', '交易日期', '股东序号'])
            await store.async_save_df(df, table_name='shareholder_top10',
                                      unique_cols=['股票代码', '交易日期', '股东序号'])
            dfs = []


async def main():
    global_dict = globals()
    tasks = []
    for key in global_dict.keys():
        if isinstance(global_dict[key], types.FunctionType) and asyncio.iscoroutinefunction(
                global_dict[key]) and key != 'main':
            tasks.append(asyncio.create_task(global_dict[key]()))

    for task in tasks:
        await task


if __name__ == '__main__':
    asyncio.run(main())
